import csv

import numpy as np
import pandas as pd
import torch
import sklearn
from sklearn.ensemble import RandomForestRegressor
from torch.utils.data import random_split

# 随机森林回归，数据样式：
# 加权销售价格,加权批发价格,是否工作日,是否节假日,春夏秋冬,分类名称_code,总销量
# 14.49,     9.23,      1,       0,       1,      0,          4.85

config = {
    'seed': 5201314,  # Your seed number, you can pick your lucky number. :)
    'select_all': True,  # Whether to use all features.
    'valid_ratio': 0.2,  # validation_size = train_size * valid_ratio
    'n_epochs': 3000,  # Number of epochs.
    'batch_size': 256,
    'learning_rate': 1e-5,
    'early_stop': 500,  # If model has not improved for this many consecutive epochs, stop training.
    'save_path': './models/model.ckpt'  # Your model will be saved here.
}


def select_feat(train_data, valid_data, test_data, select_all=True):
    '''Selects useful features to perform regression'''
    y_train, y_valid = train_data[:, -1], valid_data[:, -1]
    raw_x_train, raw_x_valid, raw_x_test = train_data[:, :-1], valid_data[:, :-1], test_data

    if select_all:
        feat_idx = list(range(raw_x_train.shape[1]))
    else:
        feat_idx = [0, 1, 2, 3, 4]  # TODO: Select suitable feature columns.

    return raw_x_train[:, feat_idx], raw_x_valid[:, feat_idx], raw_x_test[:, feat_idx], y_train, y_valid


def train_valid_split(data_set, valid_ratio, seed):
    '''Split provided training data into training set and validation set'''
    valid_set_size = int(valid_ratio * len(data_set))
    train_set_size = len(data_set) - valid_set_size
    train_set, valid_set = random_split(data_set, [train_set_size, valid_set_size],
                                        generator=torch.Generator().manual_seed(seed))
    return np.array(train_set), np.array(valid_set)


train_data, test_data = pd.read_csv('data/question2_data(1).csv', encoding="GBK").values, pd.read_csv(
    'data/question2_data(1).csv', encoding="GBK").values
train_data, valid_data = train_valid_split(train_data, config['valid_ratio'], config['seed'])

# Print out the data size.
print(f"""train_data size: {train_data.shape} 
valid_data size: {valid_data.shape} 
test_data size: {test_data.shape}""")

# Select features
xtrain, xvalid, xtest, ytrain, yvalid = select_feat(train_data, valid_data, test_data, config['select_all'])

model = RandomForestRegressor(random_state=15)
model.fit(xtrain, ytrain)
print(model.score(xvalid, yvalid))

##  加权批发价格	是否工作日	是否节假日	春夏秋冬	分类名称_code
# def pre(a, b, c, d, e):
#     #   加权销售价格	加权批发价格	是否工作日	是否节假日	春夏秋冬	分类名称_code
#     x = np.array([0, a, b, c, d, e])
#     maxProfits = 0
#     bestQuantity = 0
#     bestPrice = 0
#     for j in range(6000):
#         j = j / 1000
#         x[0] = j
#         pred = model.predict(np.array(x).reshape(1, -1))
#         profits = (j - x[1]) * pred[0]
#         if profits > maxProfits:
#             maxProfits = profits
#             bestPrice = j
#             bestQuantity = pred[0]
#     print("bestPrice :", bestPrice, "        bestQuantity :", bestQuantity, "       maxProfits :", maxProfits)
#     return [bestPrice, bestQuantity, maxProfits]
#
#
# with open("结果.csv", mode="r", encoding="GBK") as f:
#     reader = csv.reader(f)
#     # 去除首部字段名
#     header = next(reader)
#     save = [[header[7], header[8], header[9]]]
#     for row in reader:
#         temp = pre(float(row[2]), float(row[3]), float(row[4]), float(row[5]), float(row[6]))
#         save.append(temp)
#
# csvFile = open('预测.csv', 'w', encoding="GBK", newline='')
# # 如果不加入newline='',那么每写入一行数据，就会写入一行空白。可以去掉自行验证
# with csvFile:
#     writer = csv.writer(csvFile)
#     # 将array按行写入
#     writer.writerows(save)
#
